#include <openssl/evp.h>
#include <openssl/err.h>
#include <openssl/bio.h>
#include <openssl/buffer.h>
#include <openssl/rand.h>
#include <stdexcept>
#include <cstring>
#include <fstream>
#include "myAes.h"

extern "C" MyAes* createMyClass() {
    return new MyAes();
}

extern "C" void destroyMyClass(MyAes* obj) {
    delete obj;
}

MyAes::MyAes() {}
MyAes::~MyAes() {}

// 公共工具函数
std::vector<unsigned char> MyAes::base64_decode(const std::string& input) {
    BIO* b64 = BIO_new(BIO_f_base64());
    BIO* bio = BIO_new_mem_buf(input.data(), input.length());
    BIO_push(b64, bio);
    BIO_set_flags(b64, BIO_FLAGS_BASE64_NO_NL);

    std::vector<unsigned char> output(input.length());
    int len = BIO_read(b64, output.data(), input.length());
    if (len < 0) {
        BIO_free_all(b64);
        throw std::runtime_error("Base64解码失败");
    }
    output.resize(len);
    BIO_free_all(b64);
    return output;
}

std::string MyAes::base64_encode(const std::vector<unsigned char>& data) {
    BIO* b64 = BIO_new(BIO_f_base64());
    BIO* mem = BIO_new(BIO_s_mem());
    BIO_push(b64, mem);
    BIO_set_flags(b64, BIO_FLAGS_BASE64_NO_NL); // 不换行

    BIO_write(b64, data.data(), data.size());
    BIO_flush(b64);

    BUF_MEM* buf_mem;
    BIO_get_mem_ptr(b64, &buf_mem);

    std::string result(buf_mem->data, buf_mem->length);
    BIO_free_all(b64);
    return result;
}

//加密函数
std::string MyAes::openssl_aes_encrypt(const std::string& password, const std::string& plaintext) {
    ERR_load_crypto_strings();
    OpenSSL_add_all_algorithms();

    // 生成随机盐（8字节）
    unsigned char salt[8];
    if (RAND_bytes(salt, sizeof(salt)) != 1) {
        throw std::runtime_error("生成盐失败");
    }

    // 使用PBKDF2派生密钥和IV
    unsigned char key_iv[48]; // 32字节密钥 + 16字节IV
    const EVP_MD* digest = EVP_sha256();
    int iterations = 100;

    if (PKCS5_PBKDF2_HMAC(password.data(), password.length(),
                          salt, sizeof(salt),
                          iterations,
                          digest,
                          sizeof(key_iv), key_iv) != 1) {
        throw std::runtime_error("密钥派生失败");
    }

    // 初始化加密上下文
    EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
    if (!ctx) {
        throw std::runtime_error("创建加密上下文失败");
    }

    if (EVP_EncryptInit_ex(ctx, EVP_aes_256_cbc(), nullptr,
                           key_iv, key_iv + 32) != 1) {
        EVP_CIPHER_CTX_free(ctx);
        throw std::runtime_error("初始化加密失败");
    }

    // 执行加密
    std::vector<unsigned char> ciphertext(plaintext.size() + EVP_CIPHER_CTX_block_size(ctx));
    int len;
    int ciphertext_len = 0;

    if (EVP_EncryptUpdate(ctx, ciphertext.data(), &len,
                          reinterpret_cast<const unsigned char*>(plaintext.data()),
                          plaintext.size()) != 1) {
        EVP_CIPHER_CTX_free(ctx);
        throw std::runtime_error("加密失败");
    }
    ciphertext_len = len;

    if (EVP_EncryptFinal_ex(ctx, ciphertext.data() + len, &len) != 1) {
        EVP_CIPHER_CTX_free(ctx);
        throw std::runtime_error("最终加密失败");
    }
    ciphertext_len += len;
    EVP_CIPHER_CTX_free(ctx);

    // 构造最终输出数据（Salted__ + salt + ciphertext）
    std::vector<unsigned char> result;
    const char* header = "Salted__";
    result.insert(result.end(), header, header + 8);
    result.insert(result.end(), salt, salt + 8);
    result.insert(result.end(), ciphertext.begin(), ciphertext.begin() + ciphertext_len);

    // Base64编码
    std::string ciphertext_b64 = base64_encode(result);

    return ciphertext_b64;
}

// 解密函数
std::string MyAes::openssl_aes_decrypt(const std::string& password, const std::string& ciphertext_b64) {
    ERR_load_crypto_strings();
    OpenSSL_add_all_algorithms();

    // Base64解码
    std::vector<unsigned char> ciphertext = base64_decode(ciphertext_b64);

    // 验证数据格式
    if (ciphertext.size() < 16 || memcmp(ciphertext.data(), "Salted__", 8) != 0) {
        throw std::runtime_error("无效的加密格式（缺少Salted__头）");
    }

    // 提取盐（8字节）
    unsigned char salt[8];
    memcpy(salt, ciphertext.data() + 8, 8);

    // 提取实际密文
    const unsigned char* enc_data = ciphertext.data() + 16;
    int enc_len = ciphertext.size() - 16;

    // 派生密钥和IV
    unsigned char key_iv[48]; // 32字节密钥 + 16字节IV
    const EVP_MD* digest = EVP_sha256();
    int iterations = 100;

    if (PKCS5_PBKDF2_HMAC(password.data(), password.length(),
                          salt, sizeof(salt),
                          iterations,
                          digest,
                          sizeof(key_iv), key_iv) != 1) {
        throw std::runtime_error("密钥派生失败");
    }

    // 初始化解密上下文
    EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
    if (!ctx) {
        throw std::runtime_error("创建解密上下文失败");
    }

    if (EVP_DecryptInit_ex(ctx, EVP_aes_256_cbc(), nullptr,
                           key_iv, key_iv + 32) != 1) {
        EVP_CIPHER_CTX_free(ctx);
        throw std::runtime_error("初始化解密失败");
    }

    // 执行解密
    std::vector<unsigned char> plaintext(enc_len + EVP_CIPHER_CTX_block_size(ctx));
    int len;
    int plaintext_len = 0;

    if (EVP_DecryptUpdate(ctx, plaintext.data(), &len,
                          enc_data, enc_len) != 1) {
        EVP_CIPHER_CTX_free(ctx);
        throw std::runtime_error("解密失败");
    }
    plaintext_len = len;

    if (EVP_DecryptFinal_ex(ctx, plaintext.data() + len, &len) != 1) {
        EVP_CIPHER_CTX_free(ctx);
        throw std::runtime_error("最终解密失败（可能密码错误）");
    }
    plaintext_len += len;
    EVP_CIPHER_CTX_free(ctx);

    return std::string(plaintext.begin(), plaintext.begin() + plaintext_len);
}

void MyAes::openssl_aes_encrypt_file(const std::string& input_path, const std::string& output_path,const std::string& password) {
    // 读取输入文件
    std::ifstream in_file(input_path, std::ios::binary);
    if (!in_file) {
        throw std::runtime_error("无法打开输入文件: " + input_path);
    }

    // 获取文件内容
    std::string plaintext(
        (std::istreambuf_iterator<char>(in_file)),
        std::istreambuf_iterator<char>()
        );
    in_file.close();

    // 执行加密
    std::string ciphertext_b64 = openssl_aes_encrypt(password, plaintext);

    // 写入输出文件
    std::ofstream out_file(output_path);
    if (!out_file) {
        throw std::runtime_error("无法创建输出文件: " + output_path);
    }
    out_file << ciphertext_b64;
    out_file.close();
}

void MyAes::openssl_aes_decrypt_file(const std::string& input_path, const std::string& output_path,const std::string& password) {
    // 读取输入文件
    std::ifstream in_file(input_path, std::ios::binary);
    if (!in_file) {
        throw std::runtime_error("无法打开输入文件: " + input_path);
    }

    // 获取文件内容
    std::string ciphertext_b64(
        (std::istreambuf_iterator<char>(in_file)),
        std::istreambuf_iterator<char>()
        );
    in_file.close();

    // 执行解密
    std::string plaintext = openssl_aes_decrypt(password, ciphertext_b64);

    // 写入输出文件
    std::ofstream out_file(output_path);
    if (!out_file) {
        throw std::runtime_error("无法创建输出文件: " + output_path);
    }
    out_file << plaintext;
    out_file.close();
}
